import threading
from typing import Dict

from ai_platform.query_engine.rag_engine import RAGQueryEngine


class RAGEngineCache:
    """线程安全的RAGQueryEngine缓存类"""

    def __init__(self):
        self._cache: Dict[str, RAGQueryEngine] = {}
        self._lock = threading.RLock()

    def get_engine(self, session_id: str) -> RAGQueryEngine:
        """获取或创建RAGQueryEngine对象"""
        with self._lock:
            if session_id not in self._cache:
                self._cache[session_id] = RAGQueryEngine(session_id=session_id)
            return self._cache[session_id]

    def remove_engine(self, session_id: str) -> bool:
        """移除指定session的RAGQueryEngine对象"""
        with self._lock:
            if session_id in self._cache:
                del self._cache[session_id]
                return True
            return False

    def clear_cache(self):
        """清空所有缓存的RAGQueryEngine对象"""
        with self._lock:
            self._cache.clear()

    def get_cache_size(self) -> int:
        """获取当前缓存的RAGQueryEngine对象数量"""
        with self._lock:
            return len(self._cache)


# 全局RAGQueryEngine缓存实例
rag_engine_cache = RAGEngineCache()


class ChatInstance:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.query_engine = rag_engine_cache.get_engine(session_id)

    async def query(self, query: str, knowledge_base_ids: list = None):
        """
        执行查询，返回查询结果
        """
        # 获取知识库的索引类型
        index_type = "vector"  # 默认类型
        if knowledge_base_ids and len(knowledge_base_ids) > 0:
            # 获取第一个知识库的索引类型（如果有多个知识库，使用第一个的设置）
            from ai_platform.services.knowledge_service import knowledge_service
            kb = await knowledge_service.get_knowledge_base(knowledge_base_ids[0])
            if kb and kb.settings:
                index_type = kb.settings.get("index_type", "vector")

        response = await self.query_engine.query(query, knowledge_base_ids=knowledge_base_ids, index_type=index_type)
        return response

    @classmethod
    def clear_session_cache(cls, session_id: str) -> bool:
        """清除指定session的RAGQueryEngine缓存"""
        return rag_engine_cache.remove_engine(session_id)

    @classmethod
    def clear_all_cache(cls):
        """清空所有RAGQueryEngine缓存"""
        rag_engine_cache.clear_cache()

    @classmethod
    def get_cache_info(cls) -> dict:
        """获取缓存信息"""
        return {
            "cached_sessions": rag_engine_cache.get_cache_size(),
            "cache_type": "thread_safe_dict",
            "cached_object": "RAGQueryEngine"
        }

    async def get_chat_history(self, limit: int, offset: int):
        return await self.query_engine.get_chat_history(limit, offset)
